import os
import numpy as np

from net.LeNet import LeNet5, LeNet5_padding
from net.utils import *
from dataset.preprocess import *
from pipecoco.partitioner import *

from mindspore.nn import SoftmaxCrossEntropyWithLogits
from mindspore.nn import Accuracy
from mindspore import nn
from mindspore import Tensor, Model
from mindspore import context
import mindspore

context.set_context(mode=context.GRAPH_MODE, device_target='CPU')

def cal_shape(c, h, w, fused_layers):

    output_shape = []
    for j, layer in enumerate(fused_layers):
        if hasattr(layer, 'pad_mode'):
            if layer.pad_mode.lower() == 'valid':
                h,w = next_layer_size(h,w,layer.kernel_size,layer.stride)
        if hasattr(layer,'out_channels'):
            c = layer.out_channels
        output_shape.append([c,h,w])

    return output_shape

def PipeCoCo_Forward(input, fused_layers, input_shapes):

    output = [input]
    m = input_shapes[0]

    for j, input_shape in enumerate(input_shapes):
        c,h,w = input_shape
        output.append(np.zeros([m,c,h,w]))

    area_map, pre_area_map =  fused_layer_area_partitioner([5,5], 2, fused_layers[::-1])

    for i, block in enumerate(area_map.values()):
        for j, layer in enumerate(fused_layers):
            pre_layer_area = pre_area_map[i]
            x = Tensor(output[j][:, :, pre_layer_area[j][0][0]:pre_layer_area[j][0][1],
                       pre_layer_area[j][1][0]:pre_layer_area[j][1][1]], mindspore.float32)
            if j % 2 == 0 and j > 0:
                x = layers_list[-1](x)
            x = layer(x)
            k = layer.kernel_size
            if isinstance(k,tuple):
                k = k[0]
            p = 0
            if hasattr(layer, 'pad_mode'):
                if layer.pad_mode.lower() == 'same':
                    p = (k-1)//2
            output[j + 1][:, :, p+block[j][0][0]:p+block[j][0][1],
            p+block[j][1][0]:p+block[j][1][1]] = x.asnumpy()

    x = Tensor(output[-1], mindspore.float32)

    return x

if __name__=='__main__':

    mnist_path = "../datasets/MNIST_Data"
    model_path = "./models/lenet_padding/"
    network = LeNet5_padding()
    lr = 0.01
    momentum = 0.9

    # define the optimizer
    net_opt = nn.Momentum(network.trainable_params(), lr, momentum)
    # load_model(network, model_path+"/checkpoint_lenet-5_1875.ckpt")

    net_loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    model = Model(network, net_loss, metrics={"Accuracy": Accuracy()})

    ds_train = create_dataset(os.path.join(mnist_path, "train"))
    ds_eval = create_dataset(os.path.join(mnist_path, "test"))
    # acc = model.eval(ds_eval, dataset_sink_mode=False)
    # print("============== Accuracy:{} ==============".format(acc))
    train(network, net_opt, net_loss, model_path, 2, ds_train, ds_eval)
